import torch, pandas as pd
from app.services.market_crawler import CH

EDGES = [
    # 股票-所属板块
    ('sz000001', 'BK0475'), ('sz000002', 'BK0839'), ('sh600519', 'BK0730'),
    # 板块-指数
    ('BK0475', 'SZ399001'), ('BK0839', 'SZ399006'), ('BK0730', 'SH000300'),
]

def build_graph():
    df = CH.query_dataframe('''
        SELECT code, close, vol, dt
        FROM tick_1m
        WHERE dt >= now() - INTERVAL 1 DAY
    ''')
    if df.empty:
        return None, None, None
    # 节点编号
    all_nodes = list(set(df['code']) | set([n for e in EDGES for n in e]))
    node_idx = {n: i for i, n in enumerate(all_nodes)}
    edge_index = torch.tensor([[node_idx[u], node_idx[v]] for u, v in EDGES], dtype=torch.long).t().contiguous()
    # 特征：涨跌幅、量、 turnover（空数据用随机）
    feat = df.groupby('code').agg({'close': 'last', 'vol': 'sum'}).reindex(all_nodes).fillna(0)
    feat['ret'] = feat['close'].pct_change().fillna(0)
    x = torch.tensor(feat[['ret', 'vol']].values, dtype=torch.float)
    return x, edge_index, node_idx
